import sys
import os
from Bio import SeqIO
from Bio.SeqUtils import GC
from Bio.Align.Applications import ClustalwCommandline
from Bio import AlignIO

def get_filename(f):
    fileName, fileExtension = os.path.splitext(os.path.basename(f))
    return fileName


def get_extension(f):
    fileName, fileExtension = os.path.splitext(f)
    return fileExtension


def gc(record):
    return GC(record.seq)


def to_rna(record):
    return record.seq.transcribe()


def to_protein(record):
    return record.seq.translate()


def get_records(file_list):
    records = []
    for f in file_list:
        if not os.path.isfile(f):
            raise Exception ("Not a file\t" + f)
        f_type = get_extension(f)
        f_name = get_filename(f)
        if f_type == ".fasta" or f_type == ".fa":
            records.append((f_name, SeqIO.read(f, "fasta")))
        elif f_type == ".gb" or f_type == ".gbk":
            records.append((f_name, SeqIO.read(f, "genbank")))
        else:
            raise Exception ("Unknown file format\t" + f_type)
    return records


def align(proteins):
    protein_file_name = "proteins.fasta"
    output_file = open(protein_file_name, 'w')
    num = 0
    for (name, prot) in proteins:
        output_file.write(">" + name + "\n" + prot + "\n")
        num += 1
    output_file.close()
    aligned_protein_file_clustal = "aligned.clustal"
    cline = ClustalwCommandline("clustalw", infile = protein_file_name, matrix = 'BLOSUM', outfile = aligned_protein_file_clustal)
    cline()
    os.remove(protein_file_name)
    os.remove(get_filename(protein_file_name) + ".dnd")
    alignment = AlignIO.read(aligned_protein_file_clustal, "clustal")
    os.remove(aligned_protein_file_clustal)
    return alignment


def save_as_phy(alignment, phylip_file_name):
    output_handle = open(phylip_file_name, "w")
    AlignIO.write(alignment, output_handle, "phylip")
    output_handle.close()


if __name__ == '__main__':
    records = get_records(sys.argv[1:])
    proteins = []
    for (name, rec) in records:
        print "Processing ", name
        print "GC\t", gc(rec), "\n"
        print "RNA\n", to_rna(rec), "\n"
        prot = str(to_protein(rec))
        print "Protein\n", prot, "\n"
        proteins.append((name, prot))
    alignment = align(proteins)
    print "Alignment\n", alignment
    save_as_phy(alignment, "alignment.phy")

